import os
os.chdir(os.environ["HOME"])

def test_mesure():
    import gmatch4py as gm
    import networkx as nx
    import time
    from tqdm import tqdm
    import pandas as pd


    max_=100
    size_g=10
    graphs_all=[nx.random_tree(size_g) for i in range(max_)]
    result_compiled=[]
    for size_ in tqdm(range(50,max_,50)):
        graphs=graphs_all[:size_]
        comparator=None
        for class_ in [gm.BagOfNodes,gm.WeisfeleirLehmanKernel, gm.GraphEditDistance,  gm.GreedyEditDistance, gm.HED, gm.BP_2,  gm.Jaccard, gm.MCS, gm.VertexEdgeOverlap]:
            deb=time.time()
            if class_ in (gm.GraphEditDistance, gm.BP_2, gm.GreedyEditDistance, gm.HED):
                comparator = class_(1, 1, 1, 1)
            elif class_ == gm.WeisfeleirLehmanKernel:
                comparator = class_(h=2)
            else:
                comparator=class_()
            matrix = comparator.compare(graphs,None)
            print([class_.__name__,size_,time.time()-deb])
            result_compiled.append([class_.__name__,size_,time.time()-deb])

    df = pd.DataFrame(result_compiled,columns="algorithm size_data time_exec_s".split())
    df.to_csv("new_gmatch4py_res_{0}graphs_{1}size.csv".format(max_,size_g))



def MS_subgraph(mesh1, mesh2):                                                                                              
    mesh1.add_attribute("vertex_mean_curvature")          
    mesh2.add_attribute("vertex_mean_curvature")                                                                     
    radius = 0.1                                      
    curva1 = torch.from_numpy(mesh1.get_attribute("vertex_mean_curvature")).to(torch.float)                          
    curva2 = torch.from_numpy(mesh2.get_attribute("vertex_mean_curvature")).to(torch.float)                          
    ve1 =torch.from_numpy(mesh1.vertices).to(torch.float)
    curva2 = torch.from_numpy(mesh2.get_attribute("vertex_mean_curvature")).to(torch.float)          
    ve1 =torch.from_numpy(mesh1.vertices).to(torch.float)                                                            
    ve2 =torch.from_numpy(mesh2.vertices).to(torch.float)                                                            
                                                                                                                     
    # dist1                                                                                                          
    dist1 = square_distance(ve1.view(1,ve1.size()[0],ve1.size()[1]), ve1.view(1,ve1.size()[0],ve1.size()[1]))        
    N1 = ve1.size()[0]                                    
    group_idx1 = torch.arange(N1, dtype=torch.long).view(1, 1, N1).repeat([1, N1, 1])                                
    group_idx1[dist1 > radius ** 2] = N1                  
    group_idx1 = group_idx1.sort(dim=-1)[0][:, :, :]                                                                 
    group_idx1 = group_idx1.view(group_idx1.size()[1],group_idx1.size()[2])                                          
                                                                                                                     
    # dist2                                                                                                          
    dist2 = square_distance(ve2.view(1,ve2.size()[0],ve2.size()[1]), ve2.view(1,ve2.size()[0],ve2.size()[1]))        
    N2 = ve2.size()[0]                                                                                               
    # print(N1,N2)                                                                                                   
    group_idx2 = torch.arange(N2, dtype=torch.long).view(1, 1, N2).repeat([1, N2, 1])                                
    group_idx2[dist2 > radius ** 2] = N2                                                                             
    group_idx2 = group_idx2.sort(dim=-1)[0][:, :, :]      
    group_idx2 = group_idx2.view(group_idx2.size()[1],group_idx2.size()[2])                                          
                                                                                                                     
    # for loop                    
    ms_subgraph = 0                                              
    for i in range(N1):                                                                                              
        n1_idx = group_idx1[i]                
        s1_idx = n1_idx.lt(N1)                  
        f1_idx = n1_idx[s1_idx]
        g1 = mesh_to_sub_graph(mesh1,ve1.numpy(),f1_idx)                  
                                                          
        n2_idx = group_idx2[i]                            
        s2_idx = n2_idx.lt(N2)                            
        f2_idx = n2_idx[s2_idx]
        g2 = mesh_to_sub_graph(mesh2,ve2.numpy(),f2_idx)                   
        # compare                                                  
        lms_subgraph = ged.compare([g1,g2],None)   
        ms_subgraph += lms_subgraph[1][1]                                                                                  
    return ms_subgraph/N1



def mesh_to_sub_graph(mesh, vertices, v_idx):
    mesh.enable_connectivity()
    edges = []
    # Create graph
    g = nx.Graph()
    for vi in range(v_idx):
        adj_vertices = mesh.get_vertex_adjacent_vertices(vi);
        g.add_node(vi,point=vertices[vi])
        for vj in adj_vertices:
            if vj > vi:
                line = math.sqrt(pow((vertices[vi][0]-vertices[vj][0]),2)+pow((vertices[vi][1]-vertices[vj][1]),2)+pow((vertices[vi][2]-vertices[vj][2]),2))
                g.add_edge(vi,vj,weight=line)
    return g




def MSDM2(mesh1, mesh2):                                                                                              
    mesh1.add_attribute("vertex_mean_curvature")          
    mesh2.add_attribute("vertex_mean_curvature")                                                                     
    radius = 0.1                                      
    curva1 = torch.from_numpy(mesh1.get_attribute("vertex_mean_curvature")).to(torch.float)                          
    curva2 = torch.from_numpy(mesh2.get_attribute("vertex_mean_curvature")).to(torch.float)                          
    ve1 =torch.from_numpy(mesh1.vertices).to(torch.float)
    curva2 = torch.from_numpy(mesh2.get_attribute("vertex_mean_curvature")).to(torch.float)          
    ve1 =torch.from_numpy(mesh1.vertices).to(torch.float)                                                            
    ve2 =torch.from_numpy(mesh2.vertices).to(torch.float)                                                            
                                                                                                                     
    # dist1                                                                                                          
    dist1 = square_distance(ve1.view(1,ve1.size()[0],ve1.size()[1]), ve1.view(1,ve1.size()[0],ve1.size()[1]))        
    N1 = ve1.size()[0]                                    
    group_idx1 = torch.arange(N1, dtype=torch.long).view(1, 1, N1).repeat([1, N1, 1])                                
    group_idx1[dist1 > radius ** 2] = N1                  
    group_idx1 = group_idx1.sort(dim=-1)[0][:, :, :]                                                                 
    group_idx1 = group_idx1.view(group_idx1.size()[1],group_idx1.size()[2])                                          
                                                                                                                     
    # dist2                                                                                                          
    dist2 = square_distance(ve2.view(1,ve2.size()[0],ve2.size()[1]), ve2.view(1,ve2.size()[0],ve2.size()[1]))        
    N2 = ve2.size()[0]                                                                                               
    # print(N1,N2)                                                                                                   
    group_idx2 = torch.arange(N2, dtype=torch.long).view(1, 1, N2).repeat([1, N2, 1])                                
    group_idx2[dist2 > radius ** 2] = N2                                                                             
    group_idx2 = group_idx2.sort(dim=-1)[0][:, :, :]      
    group_idx2 = group_idx2.view(group_idx2.size()[1],group_idx2.size()[2])                                          
                                                                                                                     
    # for loop                    
    msdm = 0                                              
    for i in range(N1):                                                                                              
        n1_idx = group_idx1[i]                
        s1_idx = n1_idx.lt(N1)                  
        f1_idx = n1_idx[s1_idx]                           
        need_cur1 = curva1[f1_idx]                    
                                                          
        n2_idx = group_idx2[i]                            
        s2_idx = n2_idx.lt(N2)                            
        f2_idx = n2_idx[s2_idx]                                                                                      
        need_cur2 = curva2[f2_idx]                        
                                                          
        mu_1, mu_2 = need_cur1.mean(), need_cur2.mean()                                                              
        fa_1, fa_2 = need_cur1.std(), need_cur2.std()                                                                
        fa_12 = 0                                       
        fa_21 = 0                                         
        for j in f1_idx:                                                                                                     
            fa_12 += (curva1[j]-mu_1)*(curva2[j]-mu_2)    
        for q in f2_idx:                                                                                             
            fa_21 += (curva1[q]-mu_1)*(curva2[q]-mu_2)                                                                       
        fa_12 = ((fa_12/len(f1_idx))+(fa_21/len(f2_idx)))/2                                                          
        l_12 = (mu_1-mu_2).pow(2).sqrt()/(max(mu_1,mu_2)) #+1e-6)
        fa_12 = ((fa_12/len(f1_idx))+(fa_21/len(f2_idx)))/2 
        l_12 = (mu_1-mu_2).pow(2).sqrt()/(max(mu_1,mu_2)) #+1e-6)                                                    
        c_12 = (fa_1-fa_2).pow(2).sqrt()/(max(fa_1,fa_2)) #+1e-6)                                                    
        s_12 = (fa_1*fa_2-fa_12).pow(2).sqrt()/(fa_1*fa_2) #+1e-6)                                                   
        lmsdm = (0.4*l_12.pow(3) + 0.4*c_12.pow(3) + 0.2*s_12.pow(3)).pow(1/3)                                       
        if torch.isnan(lmsdm):                                                                                       
           msdm += 0                                      
        else:                                                                                                        
           print("the lmsdm pow 3 is ",lmsdm.pow(3))      
           msdm += lmsdm.pow(3)                                                                                      
    return (msdm/N1).pow(1/3)


import os
from pytorch3d.io import load_obj, save_obj
from pytorch3d.structures import Meshes
from pytorch3d.utils import ico_sphere
from pytorch3d.ops import sample_points_from_meshes
from pytorch3d.loss import (
    chamfer_distance,
    mesh_edge_loss,
    mesh_normal_consistency,
    mesh_laplacian_smoothing
)
import numpy as np
import torch
​
import matplotlib.pyplot as plt
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
from alfred.dl.torch.common import device
import open3d as o3d


trg_obj = os.path.join('./data', '10014_dolphin_v2_max2011_it2.obj')
​
verts, faces, aux = load_obj(trg_obj)


face_idx = faces.verts_idx.to(device)
verts = verts.to(device)
​
center = verts.mean(0)
verts = verts - center
scale = max(verts.abs().max(0)[0])
verts = verts / scale
​
trg_mesh = Meshes(verts=[verts], faces=[face_idx])
src_mesh = ico_sphere(4, device)

dists, idxs = _C.point_face_dist_forward( 












